package main

import (
	"flag"
	"fmt"
	"image/jpeg"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/corona10/goimagehash"
	"github.com/gogf/gf/v2/container/garray"
	"github.com/gogf/gf/v2/os/gfile"
	"github.com/gogf/gf/v2/util/gconv"
)

var (
	from     = flag.String("from", "./old", "请输入要比较的文件来源")
	find     = flag.String("find", "./find", "请输入作为源文件的目录，将从from中寻找与find目录中相似的图片")
	foundDir = flag.String("found", "./found", "输入存储找到文件的路径")
	limit    = flag.Int("limit", 8, "请输入相识度，0-99，值越小匹配度越高")
)

func main() {
	flag.Parse()
	var lockN = 60
	var (
		n        = 0
		syLock   = sync.RWMutex{}
		saveFile = "hash.log"
	)
	save, e := gfile.OpenWithFlag(saveFile, os.O_APPEND|os.O_CREATE|os.O_RDWR)
	if e != nil {
		return
	}
	defer save.Close()
	var oldMap = make(map[string]*goimagehash.ImageHash)
	// syLock.Lock()
	gfile.ReadLines(saveFile, func(line string) error {
		ss := strings.Split(line, "\t")
		if len(ss) > 1 {
			// fmt.Printf("%s\t%s\n", ss[0], ss[1])
			oldMap[ss[0]] = goimagehash.NewImageHash(gconv.Uint64(ss[1]), goimagehash.AHash)
		}
		return nil
	})
	// syLock.Unlock()
	gfile.ScanDirFileFunc(*from, "*.jpg", true, func(path string) string {
		if !strings.HasSuffix(strings.ToLower(path), ".jpg") {
			return path
		}
		if _, ok := oldMap[path]; ok {
			return path
		}
		for n >= lockN {
			time.Sleep(1 * time.Microsecond)
		}
		fmt.Printf("%s \t n=%d \n", path, n)
		go (func() {
			syLock.Lock()
			n++
			syLock.Unlock()
			defer (func() {
				syLock.Lock()
				n--
				syLock.Unlock()
			})()
			fs, e := gfile.OpenWithFlag(path, os.O_RDONLY)
			if e != nil {
				return
			}
			defer fs.Close()
			img, e := jpeg.Decode(fs)
			if e != nil {
				return
			}
			ha1, e := goimagehash.AverageHash(img)
			if e != nil {
				return
			}
			syLock.Lock()
			oldMap[path] = ha1
			save.WriteString(fmt.Sprintf("%s\t%d\n", path, ha1.GetHash()))
			syLock.Unlock()
		})()
		return path
	})
	for n > 0 {
		time.Sleep(1 * time.Second)
		fmt.Println("等待图片处理完成")
	}
	gfile.ScanDirFileFunc(*find, "*.jpg", true, func(path string) string {
		// fmt.Println(path)
		fs, e := gfile.OpenWithFlag(path, os.O_RDONLY)
		if e != nil {
			return path
		}
		defer fs.Close()
		img, e := jpeg.Decode(fs)
		if e != nil {
			return path
		}
		ha1, e := goimagehash.AverageHash(img)
		if e != nil {
			return path
		}
		// i := 1
		//只取相识度最高的3个图，如果有相识度=0的则只取相识度=0的
		limits := make(map[string]int)
		sort := garray.NewSortedIntArray()

		// limitn := 12
		for k, v := range oldMap {
			n, _ := ha1.Distance(v)
			if n < *limit {
				limits[k] = n
				sort.Add(n)
			}
		}

		sort = sort.Sort()
		slen := sort.Len()
		if slen > 3 {
			slen = 3
		}
		found := false
		for i := 0; i < slen; i++ {
			if found {
				break
			}
			if n, ok := sort.Get(i); ok {
				for k, v := range limits {
					if v == n {
						//完全相同，不用再继续查找
						if n == 0 {
							found = true
						}
						fmt.Printf("%d\t%s\n", n, path)
						gfile.CopyFile(k, gfile.Join(*foundDir, gfile.Name(k)+fmt.Sprintf("_%d_%d.%s", n, i, gfile.Ext(k))))
					}
				}
			}
		}
		return path
	})
}
